#ifndef CKKSNTT_CONTEXT_H_
#define CKKSNTT_CONTEXT_H_
#include "ciphertext.h"
#include "common.h"
#include "key.h"
#include "numb.h"
#include <NTL/ZZ.h>
#include <algorithm>
#include <arm_neon.h>
#include <chrono>
#include <cstdint>
#include <fstream>
#include <map>
#include <sstream>
#include <random>

using namespace std;
using namespace NTL;

/* *
 * @brief BFV/BGV全同态加密方案的上下文类
 * 包含各类数学运算所需的参数和基础运算函数，
 * 如NTT变换、模数转换、小工具乘法等核心操作
 *
 * @details 该类存储并管理:
 * - 环参数(N,L,K等)
 * - 模数集合(Q基,P基,T基)
 * - NTT相关参数(原根,逆元等)
 * - 基础密码学运算函数
 */

enum class params { logN15QP881 = 0, logN12QP884 = 1};

class Context {
public:
    long logN;
    long logNh;
    long L;
    long K;
    long N;
    long M;
    long Nh;
    long h;
    long T;
    int logp;
    double sigma;
    static uint64_t BASECONVERT;
    uint64_t *ex;
    uint64_t *vx;
    uint64_t *axt;
    uint64_t *bxt;
    uint64_t *axbx1;
    uint64_t *axbx2;
    uint64_t *axax;
    uint64_t *bxbx;
    uint64_t *tmpAx;
    uint64_t *tmpBx;
    uint64_t *cwQ;
    uint64_t *cwP;
    uint64_t *cwNTT;
    uint64_t **cwQSet;
    uint64_t **cwPSet;
    uint64_t *polAx;
    uint64_t *polBx;
    uint64_t *cxInvNTT;
    uint64_t *tmpCtAx;
    uint64_t *tmpCtBx;
    uint64_t *tmpP2Ql;
    uint64_t *c2QP;
    uint64_t *tmpQl2QlP;
    uint64_t *b;
    uint64_t *tmpQ2T;

    vector<uint64_t> qVec;
    vector<uint64_t> pVec;
    vector<uint64_t> tVec;
    uint64_t *twoqVec;
    uint64_t *twopVec;
    uint64_t *twotVec;
    __uint128_t *qrVec;
    __uint128_t *prVec;
    __uint128_t *trVec;
    uint64_t *qdVec;
    uint64_t *pdVec;
    uint64_t *tdVec;
    uint64_t *qrVec_high;
    uint64_t *qrVec_low;
    uint64_t *prVec_high;
    uint64_t *prVec_low;
    uint64_t *trVec_high;
    uint64_t *trVec_low;
    uint64_t *qInvVec;
    uint64_t *pInvVec;
    uint64_t *tInvVec;
    uint64_t *qRoots;
    uint64_t *pRoots;
    uint64_t *tRoots;
    uint64_t *qRootsInv;
    uint64_t *pRootsInv;
    uint64_t *tRootsInv;
    uint64_t **qRootPows;
    uint64_t **pRootPows;
    uint64_t **tRootPows;
    uint64_t **qRootPows_shoup;
    uint64_t **qRootPowsInv_shoup;
    uint64_t **pRootPows_shoup;
    uint64_t **pRootPowsInv_shoup;
    uint64_t **tRootPows_shoup;
    uint64_t **tRootPowsInv_shoup;
    uint64_t **qRootPowsInv;
    uint64_t **pRootPowsInv;
    uint64_t **tRootPowsInv;
    uint64_t *NInvModq;
    uint64_t *NInvModp;
    uint64_t *NInvModt;
    uint64_t *NInvModq_shoup;
    uint64_t *NInvModp_shoup;
    uint64_t *NInvModt_shoup;
    uint64_t **qHatModq;
    uint64_t **qHatInvModq;
    uint64_t *PModq;
    uint64_t **QModp;
    uint64_t **qInvModq;
    long *rotGroup;
    uint64_t gaussNoise = 65537;
    uint64_t *indexMatrix;
    const int GaloisGen = 5;
    ZZ *qProduct;
    ZZ *t0InvModQ;
    ZZ *pModulusAtLevel;
    uint64_t ***qHatModt;
    uint64_t **QModt;
    uint64_t **qHatInvModqShoup;
    int pow2Base;
    uint64_t **matrixValues;
    uint64_t ****qInvModqWithQP;
    uint64_t *****qInvModqpWithQP;
    uint64_t ****QModqpWithQP;
    uint64_t **pSliceHatModp;
    uint64_t **pSliceHatInvModp;
    uint64_t ***pSliceHatModq;
    uint64_t **PSliceModq;
    uint64_t **PSliceInvModq;
    uint64_t **rescaleConstants;

    Context(long logN, long L, long K, long T, int lambda = 0, int dNum = 0, long h = 64, double sigma = 3.2);

    static Context InitContext(long logN, int lambda = 0, int dNum = 0);

    void GetParameters(params param);
    void GetParametersLambda(long logN, int lambda, int dNum);

    void GetPrimes(uint64_t lambda, uint64_t logN, uint64_t bit);


    uint64_t FastPow(uint64_t a, uint64_t k, uint64_t mod);
    void QiNTT(uint64_t *res, uint64_t *a, long index);
    void PiNTT(uint64_t *res, uint64_t *a, long index);
    void NTT(uint64_t *res, uint64_t *a, long l, long k = 0);
    void QiNTTAndEqual(uint64_t *a, long index);
    void PiNTTAndEqual(uint64_t *a, long index);
    void TiNTTAndEqual(uint64_t *a, long index);
    void NTTAndEqual(uint64_t *a, long l, long k = 0);

    void QiINTT(uint64_t *res, uint64_t *a, long index);
    void PiINTT(uint64_t *res, uint64_t *a, long index);
    void INTT(uint64_t *res, uint64_t *a, long l, long k = 0);
    void QiINTTAndEqual(uint64_t *a, long index);
    void PiINTTAndEqual(uint64_t *a, long index);
    void TiINTTAndEqual(uint64_t *a, long index);
    void INTTAndEqual(uint64_t *a, long l, long k = 0);

    void QiNegate(uint64_t *res, uint64_t *a, long index);
    void PiNegate(uint64_t *res, uint64_t *a, long index);
    void Negate(uint64_t *res, uint64_t *a, long l, long k = 0);
    void QiNegateAndEqual(uint64_t *a, long index);
    void PiNegateAndEqual(uint64_t *a, long index);
    void NegateAndEqual(uint64_t *a, long l, long k = 0);
    void QiAddConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void PiAddConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void AddConst(uint64_t *res, uint64_t *a, uint64_t c, long l, long k = 0);
    void QiAddConstAndEqual(uint64_t *a, uint64_t c, long index);
    void PiAddConstAndEqual(uint64_t *a, uint64_t c, long index);
    void AddConstAndEqual(uint64_t *a, uint64_t c, long l, long k = 0);
    void QiSubConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void PiSubConst(uint64_t *res, uint64_t *a, uint64_t c, long index);
    void SubConst(uint64_t *res, uint64_t *a, uint64_t c, long l, long k = 0);
    void QiSubConstAndEqual(uint64_t *a, uint64_t c, long index);
    void PiSubConstAndEqual(uint64_t *a, uint64_t c, long index);
    void SubConstAndEqual(uint64_t *a, uint64_t c, long l, long k = 0);
    void QiAdd(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiAdd(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void TiAdd(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void Add(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void TiAddAndEqual(uint64_t *a, uint64_t *b, long index);
    void TAddAndEqual(uint64_t *a, uint64_t *b, long t);
    void AddAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiSub(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiSub(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void Sub(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);
    void TSub(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiSubAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiSubAndEqual(uint64_t *a, uint64_t *b, long index);
    void TiSubAndEqual(uint64_t *a, uint64_t *b, long index);
    void SubAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void TSubAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiSub2AndEqual(uint64_t *a, uint64_t *b, long index);
    void PiSub2AndEqual(uint64_t *a, uint64_t *b, long index);
    void Sub2AndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiMulConst(uint64_t *res, uint64_t *a, uint64_t cnst, long index);
    void PiMulConst(uint64_t *res, uint64_t *a, uint64_t cnst, long index);
    void MulConst(uint64_t *res, uint64_t *a, uint64_t cnst, long l, long k = 0);
    void QiMulConstAndEqual(uint64_t *res, uint64_t cnst, long index);
    void PiMulConstAndEqual(uint64_t *res, uint64_t cnst, long index);
    void MulConstAndEqual(uint64_t *a, uint64_t cnst, long l, long k = 0);
    void QiMul(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void PiMul(uint64_t *res, uint64_t *a, uint64_t *b, long index);
    void Mul(uint64_t *res, uint64_t *a, uint64_t *b, long l, long k = 0);
    void MulKey(uint64_t *res, uint64_t *a, uint64_t *key, long l);
    void QiMulAndEqual(uint64_t *a, uint64_t *b, long index);
    void PiMulAndEqual(uint64_t *a, uint64_t *b, long index);
    void MulAndEqual(uint64_t *a, uint64_t *b, long l, long k = 0);
    void QiSquare(uint64_t *res, uint64_t *a, long index);
    void PiSquare(uint64_t *res, uint64_t *a, long index);
    void Square(uint64_t *res, uint64_t *a, long l, long k = 0);
    void QiSquareAndEqual(uint64_t *a, long index);
    void PiSquareAndEqual(uint64_t *a, long index);
    void SquareAndEqual(uint64_t *a, long l, long k = 0);
    void EvalAndEqual(uint64_t *a, long l);
    void ReScaleAndEqual(uint64_t *&a, long l);
    void ModDownAndEqual(uint64_t *&a, long l, long dl);
    uint64_t *ModDown(uint64_t *a, long l, long dl);
    void LeftRot(uint64_t *res, uint64_t *a, long l, long rotSlots);
    void LeftRotAndEqual(uint64_t *a, long l, long rotSlots);
    void MulByMonomialAndEqual(uint64_t *a, long l, long mdeg);
    void SampleGauss(uint64_t *res, long l, long k = 0);
    void SampleZO(uint64_t *res, long l, long k = 0);
    void SampleUniform(uint64_t *res, long l, long k = 0);
    void SampleHWT(uint64_t *res, long l, long k = 0);
    uint64_t BitReverse64(uint64_t index, uint64_t bitLen);
    void PermuteNTTIndex(uint64_t galEl, uint64_t *ntt_index);
    void PermuteNTTWithIndex(uint64_t *polOut, uint64_t *polIn, uint64_t *index, int level);
    void IPermuteNTTWithIndex(uint64_t *polOut, uint64_t *polIn, uint64_t *index);
    void Encode(uint64_t *a, uint64_t *v, long slots, long l);
    void Decode(uint64_t *a, uint64_t *v, long l);
    void BaseConvertQ2T(uint64_t *a, long length, uint64_t *res);
    // void GadgetProduct(int levelQ, int levelP, uint64_t *cx, SwitchKey &gadgetCiphertext, Ciphertext &ct);
    void GadgetProduct(int levelQ, int levelP, uint64_t *cx, SwitchKey &gadgetCiphertext, uint64_t *ax, uint64_t *bx);
    void GadgetProductNoModDown(int levelQ, int levelP, uint64_t *cx, SwitchKey &gadgetCiphertext, uint64_t *tmpAx,
        uint64_t *tmpBx);
    void DecomposeSingleNTT(int levelQ, int levelP, int decompRNS, uint64_t *c2NTT, uint64_t *c2InvNTT, uint64_t *c2QiQ,
        uint64_t *c2QiP);
    void DecomposeAndSplit(int levelQ, int levelP, int decompRNS, uint64_t *p0Q, uint64_t *p1Q, uint64_t *p1P);
    void SlicesBaseConvertQl2QlP(uint64_t *a, uint64_t *QHalfModqi, uint64_t *p1Q, uint64_t *p1P, int levelQ,
        int levelP, int decompLvl, long length, long slices, long offset, ZZ QHalf);
    void GadgetProductSinglePAndBitDecompNoModDown(int levelQ, int levelP, uint64_t *cx, SwitchKey &gadgetCt,
        uint64_t *tmpAx, uint64_t *tmpBx);
    void BaseConvertQlPk2Ql(uint64_t *res, uint64_t *a, long levelQ, long levelP);
    void ExtendBasisSmallNormAndCenter(uint64_t *res, uint64_t *a, int levelP);
};
#endif